/**
 * @brief 使用libevauth自定义推理 yolov5s
 */

#include "evdeploy/cv/detection/yolov5_prepostprocessor.h"
#include "model_handle.h"
#include "opencv2/core.hpp"
#include "opencv2/opencv.hpp"
#include <iostream>
#include <memory>

using namespace ev;
int main(int argc, char **argv)
{
    const std::string onnx_path = argv[1]; // yolov5sonnx模型地址
    const std::string pic_path = argv[2];  //输入图片的路径
    printf("model:%s image:%s\n", onnx_path.c_str(), pic_path.c_str());

    // 转换模型与加载模型
    auto engine = std::make_shared<EVCudaEngine>();
    std::string trt_path = onnx_path.substr(0, onnx_path.find_last_of(".")) + ".trt";
    int ret;
    if (CheckFileType(trt_path.c_str()) != EV_FILE)
    {
        ret = engine->convertEngine(onnx_path);
        if (ret != 0)
        {
            std::cout << "model convert failed" << std::endl;
            exit(-1);
        }
    }
    printf("loading trt:%s\n", trt_path.c_str());

    // 加载模型
    ret = engine->loadModel(trt_path);
    if (ret != 0)
    {
        std::cout << "model convert failed" << std::endl;
        exit(-1);
    }
    auto context = engine->createExecutionContext();

    // 读取图片
    cv::Mat src = cv::imread(pic_path);
    if (src.empty())
    {
        std::cout << "pic is empty" << std::endl;
        exit(-1);
    }

    // 前处理
    ev::vision::YOLOv5Preprocessor preprocessor;
    ev::vision::YOLOv5Postprocessor postprocessor;
    cv::Mat cv_in_mat1;
    preprocessor.Run(src, cv_in_mat1, 640);

    // init buffer
    void *buffers[2];
    EVMatData ev_data;
    const int inputIndex = engine->getBindingIndex("images");
    const int outputIndex = engine->getBindingIndex("output");
    auto in_dims = engine->getBindingDimensions(inputIndex);
    auto out_dims = engine->getBindingDimensions(outputIndex);
    // const int input_size = 640 * 640 * 3 * sizeof(float);
    int input_size = 1;
    for (int i = 0; i < in_dims.nbDims; i++)
    {
        input_size *= in_dims.d[i];
        std::cout << in_dims.d[i] << " ";
    }
    std::cout << " in tatal:" << input_size << std::endl;

    int output_size = 1;
    for (int i = 0; i < out_dims.nbDims; i++)
    {
        output_size *= out_dims.d[i];
        ev_data.dims.push_back(out_dims.d[i]);
        std::cout << out_dims.d[i] << " ";
    }
    std::cout << " out tatal:" << output_size << std::endl;
    std::vector<float> output_buffer(output_size, 0);
    cudaMalloc(&buffers[inputIndex], input_size * sizeof(float));
    cudaMalloc(&buffers[outputIndex], output_size * sizeof(float));

    // 推理
    cudaStream_t stream;
    cudaStreamCreate(&stream);
    cudaMemcpyAsync(buffers[inputIndex], cv_in_mat1.data, input_size * sizeof(float), cudaMemcpyHostToDevice, stream);
    context->enqueueV2(buffers, stream, nullptr);
    cudaMemcpyAsync(output_buffer.data(), buffers[outputIndex], output_size * sizeof(float), cudaMemcpyDeviceToHost,
                    stream);
    cudaStreamSynchronize(stream);
    cudaStreamDestroy(stream);
    cudaFree(buffers[inputIndex]);
    cudaFree(buffers[outputIndex]);

    // 后处理
    ev_data.data = output_buffer.data();
    std::vector<ev::vision::BoxInfo> objects;
    postprocessor.Run(&ev_data, objects, preprocessor.GetScale(), 0.5, src.cols, src.rows);
    for (const auto &obj : objects)
    {
        printf("box:[%f,%f,%f,%f] conf:%0.2f label:%d\n", obj.x1, obj.y1, obj.x2, obj.y2, obj.score, obj.label);
        cv::Rect rect(cv::Point(int(obj.x1), int(obj.y1)), cv::Point(int(obj.x2), int(obj.y2)));
        cv::rectangle(src, rect, cv::Scalar(0, 0, 255));
    }
    cv::imwrite("ret.jpg", src);
    return 0;
}